package com.katze.boot.applet.ssh;

import com.katze.boot.applet.ssh.service.SecureShellService;
import com.katze.boot.file.entity.FileTransfer;
import com.katze.boot.file.servlet.DefaultFileTransferServlet;
import com.katze.boot.file.servlet.IFileUploadService;
import com.katze.ssh.SshSessionPool;
import com.katze.ssh.sftp.SftpChannel;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import org.apache.commons.lang3.StringUtils;
import org.apache.sshd.sftp.client.SftpClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.web.servlet.ServletRegistrationBean;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.channels.FileChannel;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Stream;

@Component
public class SftpFileUploadService implements IFileUploadService {
    private static final Logger log = LoggerFactory.getLogger(SftpFileUploadService.class);

    @Resource
    private SshSessionPool sessionPool;
    @Resource
    private SecureShellService secureShellService;
    @Resource
    private ServletRegistrationBean<DefaultFileTransferServlet> servletRegistrationBean;

    @PostConstruct
    public void init() {
        servletRegistrationBean.getServlet().add("sftp", this);
    }

    @Override
    public String[] beforeUpload(FileTransfer transfer) {
        File tempDir = new File(transfer.getTemp(), transfer.getOwner() + "_" + transfer.getId());
        if (!tempDir.exists() && !tempDir.mkdirs()) {
            log.error("{}目录创建失败", tempDir);
            throw new IllegalStateException("当前上传功能暂停使用");
        }
        String[] chunks = tempDir.list((dir, name) -> StringUtils.isNumeric(name));
        return chunks == null? new String[]{}:chunks;
    }

    @Override
    public void upload(FileTransfer transfer) throws IOException {
        Path dest = transfer.getTemp().toPath().resolve((transfer.getOwner() + "_" + transfer.getId()) + "/" + transfer.getIndex());
        if (Files.exists(dest)) {
            log.warn("{}文件将被覆盖", dest);
        }
        Files.write(dest, transfer.getData(), StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
    }

    @Override
    public void afterUpload(FileTransfer transfer) {
        File dir = new File(transfer.getTemp(), transfer.getOwner() + "_" + transfer.getId());
        try (SftpChannel channel = sessionPool.borrow(secureShellService.getSessionProperties(transfer.getFrom()), SftpChannel.class)){
            File[] files = dir.listFiles();
            if (files != null && files.length > 0) {
                Stream<File> stream = Arrays.stream(files).sorted(Comparator.comparing(o -> Integer.valueOf(o.getName())));
                String tempDir = channel.mkdirs(channel.getHomePath(), "temp");
                String tempFile = tempDir + "/" + transfer.getId();
                SftpClient.Attributes attributes = channel.exists(tempFile);
                if (attributes != null) {
                    long size = new File(dir, "0").length();
                    if (files.length == 1 || (attributes.getSize() != transfer.getSize() && attributes.getSize() % size != 0)) {
                        // 1个文件块数或者文件块大小发生变化，则直接覆盖
                        channel.rm(tempDir, transfer.getId());
                        stream =  stream.filter(item -> StringUtils.isNumeric(item.getName()));
                    } else {
                        int count = (int) Math.ceil(attributes.getSize() * 1.0 / size); // 向上取整，获取已上传块数
                        // 追加
                        stream =  stream.filter(item -> Integer.parseInt(item.getName()) >= count);
                    }
                } else {
                    stream =  stream.filter(item -> StringUtils.isNumeric(item.getName()));
                }
                List<File> chunks = stream.toList();
                if (!chunks.isEmpty()) {
                    channel.put(tempFile, ch -> {
                        for (File chunk : chunks) {
                            try (FileInputStream fileInputStream = new FileInputStream(chunk); FileChannel inChannel = fileInputStream.getChannel()){
                                inChannel.transferTo(0, inChannel.size(), ch);
                            } catch (IOException e) {
                                log.error("文件块目录:{}, 实际文件名:{}, 文件序号:{}合并失败", dir, transfer.getName(), chunk.getName());
                                throw new IllegalStateException("文件合并失败");
                            }
                            if (!chunk.delete()) {
                                log.warn(chunk.getAbsolutePath() + "文件删除失败");
                            }
                        }
                    });
                }
                if (transfer.getRelativePath() != null) {
                    transfer.setTarget(channel.mkdirs(transfer.getTarget(), transfer.getRelativePath()));
                }
                channel.mv(tempDir, transfer.getId(), transfer.getTarget(), transfer.getName());
            } else {
                throw new IllegalStateException("文件已超时重新上传");
            }
        } catch (Exception e) {
            log.error("文件保存失败", e);
            throw new IllegalStateException("文件保存失败");
        } finally {
            String[] fs = dir.list();
            if (fs != null && fs.length == 0 && !dir.delete()) log.warn("{}目录删除失败", dir);
        }
    }
}
